#include "common.h"
#include <fstream>
#include <sstream>

// 将AVFrame转换为RGB24格式
AVFrame* convert_to_rgb24(AVFrame* src_frame, int width, int height) {
    AVFrame* rgb_frame = av_frame_alloc();
    if (!rgb_frame) {
        return nullptr;
    }
    
    rgb_frame->format = AV_PIX_FMT_RGB24;
    rgb_frame->width = width;
    rgb_frame->height = height;
    
    int ret = av_frame_get_buffer(rgb_frame, 32);
    if (ret < 0) {
        av_frame_free(&rgb_frame);
        return nullptr;
    }
    
    SwsContext* sws_ctx = sws_getContext(
        width, height, (AVPixelFormat)src_frame->format,
        width, height, AV_PIX_FMT_RGB24,
        SWS_BILINEAR, nullptr, nullptr, nullptr
    );
    
    if (!sws_ctx) {
        av_frame_free(&rgb_frame);
        return nullptr;
    }
    
    sws_scale(sws_ctx, src_frame->data, src_frame->linesize,
              0, height, rgb_frame->data, rgb_frame->linesize);
    
    sws_freeContext(sws_ctx);
    return rgb_frame;
}

// 保存帧为PPM格式
bool save_frame_as_ppm(AVFrame* frame, const std::string& filename) {
    std::ofstream file(filename, std::ios::binary);
    if (!file.is_open()) {
        return false;
    }
    
    file << "P6\n" << frame->width << " " << frame->height << "\n255\n";
    
    for (int y = 0; y < frame->height; y++) {
        file.write((char*)(frame->data[0] + y * frame->linesize[0]), 
                   frame->width * 3);
    }
    
    file.close();
    return true;
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        std::cerr << "使用方法: " << argv[0] << " <视频文件> <帧索引>" << std::endl;
        std::cerr << "示例: " << argv[0] << " test.mp4 100" << std::endl;
        return 1;
    }
    
    init_ffmpeg();
    
    std::string filename = argv[1];
    int target_frame = std::stoi(argv[2]);
    
    AVFormatContext* fmt_ctx = nullptr;
    int ret = avformat_open_input(&fmt_ctx, filename.c_str(), nullptr, nullptr);
    CHECK_ERROR(ret, "无法打开文件");
    
    ret = avformat_find_stream_info(fmt_ctx, nullptr);
    CHECK_ERROR(ret, "无法获取流信息");
    
    int video_stream_index = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
    if (video_stream_index < 0) {
        std::cerr << "未找到视频流" << std::endl;
        avformat_close_input(&fmt_ctx);
        return 1;
    }
    
    AVStream* video_stream = fmt_ctx->streams[video_stream_index];
    const AVCodec* codec = avcodec_find_decoder(video_stream->codecpar->codec_id);
    if (!codec) {
        std::cerr << "未找到解码器" << std::endl;
        avformat_close_input(&fmt_ctx);
        return 1;
    }
    
    AVCodecContext* codec_ctx = avcodec_alloc_context3(codec);
    if (!codec_ctx) {
        std::cerr << "无法分配解码器上下文" << std::endl;
        avformat_close_input(&fmt_ctx);
        return 1;
    }
    
    ret = avcodec_parameters_to_context(codec_ctx, video_stream->codecpar);
    CHECK_ERROR(ret, "无法复制编解码器参数");
    
    ret = avcodec_open2(codec_ctx, codec, nullptr);
    CHECK_ERROR(ret, "无法打开解码器");
    
    // 计算目标时间戳
    int64_t target_pts = (int64_t)target_frame * video_stream->time_base.den / 
                         (int64_t)(video_stream->time_base.num * av_q2d(av_guess_frame_rate(fmt_ctx, video_stream, nullptr)));
    
    // 跳转到目标帧附近
    ret = av_seek_frame(fmt_ctx, video_stream_index, target_pts, AVSEEK_FLAG_BACKWARD);
    if (ret < 0) {
        std::cerr << "无法跳转到指定帧" << std::endl;
        avcodec_free_context(&codec_ctx);
        avformat_close_input(&fmt_ctx);
        return 1;
    }
    
    AVPacket* packet = av_packet_alloc();
    AVFrame* frame = av_frame_alloc();
    int frame_count = 0;
    bool frame_found = false;
    
    while (av_read_frame(fmt_ctx, packet) >= 0) {
        if (packet->stream_index == video_stream_index) {
            ret = avcodec_send_packet(codec_ctx, packet);
            if (ret < 0) continue;
            
            ret = avcodec_receive_frame(codec_ctx, frame);
            if (ret == 0) {
                if (frame_count >= target_frame) {
                    AVFrame* rgb_frame = convert_to_rgb24(frame, codec_ctx->width, codec_ctx->height);
                    if (rgb_frame) {
                        std::stringstream filename;
                        filename << "frame_" << std::setfill('0') << std::setw(6) << target_frame << ".ppm";
                        
                        if (save_frame_as_ppm(rgb_frame, filename.str())) {
                            std::cout << "✅ 成功保存帧 " << target_frame << " 到 " << filename.str() << std::endl;
                            std::cout << "📐 尺寸: " << rgb_frame->width << " × " << rgb_frame->height << std::endl;
                        } else {
                            std::cerr << "❌ 无法保存帧" << std::endl;
                        }
                        
                        av_frame_free(&rgb_frame);
                        frame_found = true;
                    }
                    break;
                }
                frame_count++;
            }
        }
        av_packet_unref(packet);
    }
    
    if (!frame_found) {
        std::cerr << "❌ 无法找到帧 " << target_frame << std::endl;
    }
    
    av_frame_free(&frame);
    av_packet_free(&packet);
    avcodec_free_context(&codec_ctx);
    avformat_close_input(&fmt_ctx);
    
    return frame_found ? 0 : 1;
}
